from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import OneHotEncoder
import pandas as pd
import numpy as np

# Read dataset
dataset = pd.read_csv('../dataset/动物数据集.xlsx')

# Extract features and labels
x = dataset[['体型', '翅膀数量', '脚数量', '是否产蛋']].values
y = dataset['类别'].values

# Separate numerical and categorical features
numerical_features = x[:, 1:3].astype(float)  # Ensure numerical features are float
categorical_features = x[:, [0, 3]]  # Assume "体型" and "是否产蛋" are categorical features

# Perform one-hot encoding on categorical features
one_hot_encoder = OneHotEncoder(handle_unknown='ignore', sparse_output=False)
encoded_categorical_features = one_hot_encoder.fit_transform(categorical_features)

# Combine numerical features and encoded categorical features
encoded_data = np.hstack((numerical_features, encoded_categorical_features))

# Train logistic regression model
model = LogisticRegression()
model.fit(encoded_data, y)

# Predict new sample
prediction_input = np.array([['中', 0, 4, '否']])
numerical_prediction_input = prediction_input[:, 1:3].astype(float)  # Ensure numerical features are float
categorical_prediction_input = prediction_input[:, [0, 3]]  # Assume "体型" and "是否产蛋" are categorical features
encoded_prediction_input = one_hot_encoder.transform(categorical_prediction_input)
prediction_input_data = np.hstack((numerical_prediction_input, encoded_prediction_input))
prediction_result = model.predict(prediction_input_data)
print(prediction_result)
